#!/usr/bin/env python3
"""
mpc_quad_mpc_hover_vis.py

- Receding-horizon LQR-as-MPC proxy for 3D double-integrator.
- When reached target: switch to hover mode (PD controller) for hover_time seconds.
- Real-time prints with explanations of MPC output & decision logic.
- Real-time visualization:
    Left: 3D (trajectory, start, target, predicted future trajectory)
    Right top: positions vs time (with predicted future)
    Right bottom: control vs time (with predicted future controls as markers)
- Save CSV log and combined PNG on finish.

Tune parameters at the top for different ranges/behavior.
"""
import os
import time as pytime
from math import ceil
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# ---------------------------
# Parameters (tweak these)
# ---------------------------
dt = 0.1
T_total = 60.0          # max sim time (s)
steps = int(ceil(T_total / dt))
horizon = 50            # prediction horizon for K_seq
pred_show = 10          # how many predicted steps to show in visualization/prints
hover_time = 10.0       # seconds to hover after arriving
pos_tol = 0.5           # arrival distance tolerance (m)
vel_tol = 0.5           # arrival speed tolerance (m/s)

# actuation limits (m/s^2)
u_min = np.array([-5.0, -5.0, -5.0])
u_max = np.array([ 5.0,  5.0,  5.0])

# target & initial
target_pos = np.array([80.0, -20.0, 30.0])   # goal within 100m
# initial state [x,y,z, vx,vy,vz]
x0 = np.zeros(6)
x0[2] = 0.0

# model (3D double integrator discrete)
I3 = np.eye(3)
A = np.block([[I3, dt * I3],
              [np.zeros((3,3)), I3]])
B = np.block([[0.5 * dt**2 * I3],
              [dt * I3]])

# cost matrices (LQR)
Q_pos = 10.0
Q_vel = 1.0
Q = np.diag([Q_pos, Q_pos, Q_pos, Q_vel, Q_vel, Q_vel])
R = 0.1 * np.eye(3)

# disturbance (optional)
np.random.seed(1)
wind_amp = 0.05
def wind(t):
    return wind_amp * (np.random.randn(3) * 0.5)

# ---------------------------
# Helper: finite-horizon LQR gains (backward Riccati)
# ---------------------------
def finite_horizon_lqr_gains(A, B, Q, R, N):
    P = [None] * (N + 1)
    P[N] = Q.copy()
    K = [None] * N
    for k in range(N-1, -1, -1):
        Bt_PB = B.T @ P[k+1] @ B
        inv = np.linalg.inv(R + Bt_PB)
        Kk = inv @ (B.T @ P[k+1] @ A)
        K[k] = Kk
        Pk = Q + A.T @ P[k+1] @ A - A.T @ P[k+1] @ B @ Kk
        P[k] = Pk
    return K, P

K_seq, P_seq = finite_horizon_lqr_gains(A, B, Q, R, horizon)
K0 = K_seq[0]

# ---------------------------
# Prediction helper (closed-loop with K_seq)
# ---------------------------
def predict_trajectory(x0, x_ref, K_seq, A, B, steps):
    """
    Predict closed-loop states when applying u = -K_seq[k] * (x - x_ref)
    Returns preds shape (steps+1, 6) and u_preds shape (steps,3)
    """
    x = x0.copy()
    preds = [x.copy()]
    u_preds = []
    for k in range(min(steps, len(K_seq))):
        err = x - x_ref
        u = -K_seq[k].dot(err)
        preds.append((A.dot(x) + B.dot(u)).copy())
        u_preds.append(u.copy())
        x = preds[-1]
    return np.array(preds), np.array(u_preds)

# ---------------------------
# PD hover controller (used after arrival)
# ---------------------------
def pd_hover_control(pos, vel, target_pos, kp=1.0, kd=1.5):
    # simple PD on position to hold target (tuned small)
    pos_err = pos - target_pos
    u = -kp * pos_err - kd * vel
    return u

# ---------------------------
# Setup logging and plotting
# ---------------------------
out_dir = os.path.dirname(os.path.abspath(__file__))
log_path = os.path.join(out_dir, 'mpc_quad_mpc_hover_log.csv')
png_path = os.path.join(out_dir, 'mpc_quad_mpc_hover_combined.png')

times = []
x_hist = []
u_hist = []
u_raw_hist = []

# state
x = x0.copy()
x_ref = np.zeros(6); x_ref[0:3] = target_pos

# plotting - interactive with fallback
interactive_ok = True
fig = None
try:
    plt.ion()
    fig = plt.figure(figsize=(13,7))
    gs = fig.add_gridspec(2, 3, width_ratios=[2.5,1,1])
    ax3d = fig.add_subplot(gs[:,0], projection='3d')
    ax_pos = fig.add_subplot(gs[0,1])
    ax_ctrl = fig.add_subplot(gs[1,1])
    ax_info = fig.add_subplot(gs[:,2])  # use this for textual explanation (turn off axes)

    # initial plot elements
    traj_line, = ax3d.plot([], [], [], lw=2, label='trajectory')
    start_scatter = ax3d.scatter([x0[0]], [x0[1]], [x0[2]], c='b', s=80, marker='o', label='start')
    target_scatter = ax3d.scatter([target_pos[0]], [target_pos[1]], [target_pos[2]], c='r', s=120, marker='x', label='target')
    pred_line, = ax3d.plot([], [], [], ls='--', lw=1.5, color='orange', label='MPC-pred')
    pred_dots = ax3d.scatter([], [], [], c='orange', s=30)
    final_scatter = ax3d.scatter([], [], [], c='g', s=150, marker='*', label='final')

    ax3d.set_xlabel('x (m)'); ax3d.set_ylabel('y (m)'); ax3d.set_zlabel('z (m)')
    ax3d.set_title('3D Trajectory (start blue, target red, pred orange, final green)')
    ax3d.legend(loc='upper left')

    pos_lines = [ax_pos.plot([], [], label='x')[0],
                 ax_pos.plot([], [], label='y')[0],
                 ax_pos.plot([], [], label='z')[0]]
    ax_pos.set_xlabel('time (s)'); ax_pos.set_ylabel('position (m)')
    ax_pos.set_title('Position vs time'); ax_pos.legend(); ax_pos.grid(True)

    ctrl_lines = [ax_ctrl.plot([], [], label='ax')[0],
                  ax_ctrl.plot([], [], label='ay')[0],
                  ax_ctrl.plot([], [], label='az')[0]]
    ax_ctrl.set_xlabel('time (s)'); ax_ctrl.set_ylabel('acc (m/s^2)')
    ax_ctrl.set_title('Control vs time'); ax_ctrl.legend(); ax_ctrl.grid(True)

    # info axis: we will draw text inside
    ax_info.axis('off')
    info_text = ax_info.text(0.01, 0.98, "", va='top', ha='left', fontsize=9, family='monospace')

    plt.tight_layout()
except Exception as e:
    print("Interactive plotting unavailable, will save final images. Exception:", e)
    interactive_ok = False
    try:
        plt.ioff()
    except:
        pass

def update_visuals(t_arr, X_arr, U_arr, preds, u_preds, hover_flag, info_str):
    global interactive_ok, traj_line, pred_line, pred_dots, final_scatter, info_text
    if not interactive_ok or X_arr.shape[0] == 0:
        return
    try:
        # 3D real trajectory
        traj_line.set_data(X_arr[:,0], X_arr[:,1])
        traj_line.set_3d_properties(X_arr[:,2])
        # predicted line (only positions)
        if preds is not None and preds.shape[0] > 0:
            pred_line.set_data(preds[:,0], preds[:,1])
            pred_line.set_3d_properties(preds[:,2])
            # update pred dots
            pred_dots.remove()
            pred_dots = ax3d.scatter(preds[:,0], preds[:,1], preds[:,2], c='orange', s=30)
        else:
            pred_line.set_data([], []); pred_line.set_3d_properties([])
            try:
                pred_dots.remove()
            except:
                pass

        # dynamic axis limits
        margin = 10.0
        ax3d.set_xlim(np.min(X_arr[:,0]) - margin, np.max(X_arr[:,0]) + margin)
        ax3d.set_ylim(np.min(X_arr[:,1]) - margin, np.max(X_arr[:,1]) + margin)
        ax3d.set_zlim(0, max(np.max(X_arr[:,2]) + margin, target_pos[2] + margin))

        # pos & ctrl lines
        for i, ln in enumerate(pos_lines):
            ln.set_data(t_arr, X_arr[:, i])
        ax_pos.set_xlim(t_arr[0], t_arr[-1])
        all_pos = X_arr[:,0:3].flatten()
        ax_pos.set_ylim(np.min(all_pos)-5, np.max(all_pos)+5)

        for i, ln in enumerate(ctrl_lines):
            ln.set_data(t_arr, U_arr[:, i])
        ax_ctrl.set_xlim(t_arr[0], t_arr[-1])
        all_ctrl = U_arr.flatten()
        ax_ctrl.set_ylim(np.min(all_ctrl)-1.0, np.max(all_ctrl)+1.0)

        # info text (monospace)
        info_text.set_text(info_str)

        fig.canvas.draw(); fig.canvas.flush_events()
    except Exception as e:
        print("Plot update error, disabling live updates:", e)
        interactive_ok = False

# ---------------------------
# Main simulation loop
# ---------------------------
print("Starting simulation. steps =", steps, " dt =", dt, " horizon =", horizon)
sim_start = pytime.time()

hover_mode = False
hover_start_time = None
reached = False

for step in range(steps):
    t = step * dt
    err = x - x_ref
    # MPC "u_raw" (linear receding-horizon LQR proxy)
    u_raw = -K0.dot(err)
    # predict closed-loop future using K_seq for visualization/explanation
    preds_all, u_preds_all = predict_trajectory(x, x_ref, K_seq, A, B, pred_show)
    pred_positions = preds_all[:,0:3] if preds_all is not None else None

    # if already in hover_mode -> use PD controller to keep position
    if hover_mode:
        pos = x[0:3]; vel = x[3:6]
        u_des = pd_hover_control(pos, vel, target_pos, kp=0.6, kd=1.0)
        u = np.maximum(u_min, np.minimum(u_max, u_des))
        reason = "HOVER_PD"
    else:
        # apply clipping to u_raw for safety
        u = np.maximum(u_min, np.minimum(u_max, u_raw))
        reason = "MPC_CLIPPED" if not np.allclose(u, u_raw) else "MPC_OK"

    # apply disturbance
    w = wind(t)
    x = A.dot(x) + B.dot(u) + B.dot(w)

    # logs
    times.append(t); x_hist.append(x.copy()); u_hist.append(u.copy()); u_raw_hist.append(u_raw.copy())

    # metrics
    pos = x[0:3]; vel = x[3:6]
    err_pos_norm = np.linalg.norm(err[0:3])
    u_raw_norm = np.linalg.norm(u_raw); u_norm = np.linalg.norm(u)
    per_axis_clipped = (u_raw < u_min) | (u_raw > u_max)
    sat_axes = int(np.sum(per_axis_clipped))
    pred_final = pred_positions[-1] if pred_positions is not None and pred_positions.shape[0]>0 else None
    pred_final_dist = np.linalg.norm(pred_final - target_pos) if pred_final is not None else None

    # Build info string for figure
    info_lines = []
    info_lines.append(f"t={t:5.2f}s  mode={'HOVER' if hover_mode else 'MPC'}  reason={reason}")
    info_lines.append(f"pos=({pos[0]:6.1f},{pos[1]:6.1f},{pos[2]:5.1f})  vel=({vel[0]:5.1f},{vel[1]:5.1f},{vel[2]:5.1f})")
    info_lines.append(f"u_raw=({u_raw[0]:7.2f},{u_raw[1]:7.2f},{u_raw[2]:7.2f})  => u=({u[0]:5.2f},{u[1]:5.2f},{u[2]:5.2f})  ||u_raw||={u_raw_norm:7.2f} ||u||={u_norm:6.2f}")
    info_lines.append(f"sat_axes={sat_axes}/3  err_pos={err_pos_norm:6.2f}m")
    if pred_final is not None:
        info_lines.append(f"pred_final (h={pred_show}) = ({pred_final[0]:6.1f},{pred_final[1]:6.1f},{pred_final[2]:5.1f}) dist_to_target_pred={pred_final_dist:6.2f}")
    info_lines.append("")
    info_lines.append("MPC decision logic (short):")
    info_lines.append("  1) compute u_raw = -K0*(x - x_ref) (instant feedback from finite-horizon LQR).")
    info_lines.append("  2) predict future states applying u = -K_seq[k]*(x-x_ref) (we show first few predicted positions).")
    info_lines.append("  3) clip u_raw into safe actuator bounds -> actual u applied.")
    info_lines.append("  4) UAV moves according to dynamics; next step repeats.")
    info_lines.append("")
    info_lines.append("If many axes saturated (sat_axes>0) -> prediction assumes larger control than executed;")
    info_lines.append(" -> consider smoothing reference / increasing R or using constrained MPC (solver handles bounds).")
    info_str = "\n".join(info_lines)

    # Console prints (concise + explanation)
    print(f"[t={t:5.2f}] pos=({pos[0]:6.1f},{pos[1]:6.1f},{pos[2]:5.1f}) vel=({vel[0]:5.2f},{vel[1]:5.2f},{vel[2]:5.2f})")
    print(f"    u_raw=({u_raw[0]:7.2f},{u_raw[1]:7.2f},{u_raw[2]:7.2f}) -> u=({u[0]:5.2f},{u[1]:5.2f},{u[2]:5.2f}) sat_axes={sat_axes}/3 ||u_raw||={u_raw_norm:7.2f}")
    if pred_positions is not None:
        showN = min(predictions := pred_positions.shape[0]-1, pred_show)
        pred_list = " ".join([f"({pred_positions[i+1,0]:5.1f},{pred_positions[i+1,1]:5.1f},{pred_positions[i+1,2]:4.1f})" for i in range(showN)])
        print(f"    Pred next {showN}: {pred_list}  pred_final_dist={pred_final_dist:6.2f}")
    # Explain how next motion chosen
    print("    >> How MPC decides next motion: compute u_raw to reduce future Q-cost;")
    print("       If u_raw inside bounds, u_raw is applied; else clipped -> dynamics differ from prediction.")
    if sat_axes > 0:
        print("    >> DIAG: saturation detected - controller requested more than actuators allow.")

    # Arrival detection: if within tol and slow => enter hover
    if (err_pos_norm < pos_tol) and (np.linalg.norm(vel) < vel_tol) and (not hover_mode):
        hover_mode = True
        hover_start_time = t
        print(f"*** ARRIVED at target: switching to HOVER mode for {hover_time} s (t={t:.2f}). ***")

    # When in hover_mode, stop after hover_time seconds
    if hover_mode:
        if (t - hover_start_time) >= hover_time:
            print("*** HOVER complete: ending simulation. ***")
            reached = True
            break

    # update plots (real-time)
    try:
        t_arr = np.array(times + [t])
        X_arr = np.array(x_hist + [x])
        U_arr = np.array(u_hist + [u])
        update_visuals(t_arr, X_arr, U_arr, pred_positions, u_preds_all, hover_mode, info_str)
        if interactive_ok:
            plt.pause(0.001)
    except KeyboardInterrupt:
        print("Simulation interrupted by user.")
        break
    except Exception as e:
        print("Runtime plotting error:", e)
        break

sim_end = pytime.time()
print(f"Simulation ended. wall_time={sim_end-sim_start:.2f}s steps_simulated={len(times)} reached={reached}")

# Save log CSV
Xnp = np.array(x_hist)
Unp = np.array(u_hist)
df = pd.DataFrame({
    't': np.array(times),
    'x': Xnp[:,0], 'y': Xnp[:,1], 'z': Xnp[:,2],
    'vx': Xnp[:,3], 'vy': Xnp[:,4], 'vz': Xnp[:,5],
    'ax': Unp[:,0], 'ay': Unp[:,1], 'az': Unp[:,2]
})
df.to_csv(log_path, index=False)
print("Saved CSV log to:", log_path)

# Save final figure (ensure final marker)
try:
    final_pos = Xnp[-1,0:3]
    if interactive_ok and fig is not None:
        ax3d.scatter([final_pos[0]], [final_pos[1]], [final_pos[2]], c='g', s=150, marker='*', label='final')
        ax3d.legend()
        fig.savefig(png_path, dpi=160, bbox_inches='tight')
    else:
        # static compose
        figs = plt.figure(figsize=(13,7))
        gs2 = figs.add_gridspec(2,3, width_ratios=[2.5,1,1])
        ax3d_s = figs.add_subplot(gs2[:,0], projection='3d')
        ax_pos_s = figs.add_subplot(gs2[0,1]); ax_ctrl_s = figs.add_subplot(gs2[1,1]); ax_info_s = figs.add_subplot(gs2[:,2])
        ax3d_s.plot(Xnp[:,0], Xnp[:,1], Xnp[:,2], lw=2)
        ax3d_s.scatter([x0[0]],[x0[1]],[x0[2]], c='b', s=80, marker='o', label='start')
        ax3d_s.scatter([target_pos[0]],[target_pos[1]],[target_pos[2]], c='r', s=120, marker='x', label='target')
        ax3d_s.scatter([final_pos[0]],[final_pos[1]],[final_pos[2]], c='g', s=150, marker='*', label='final')
        ax3d_s.set_title('3D Trajectory (final)')
        ax3d_s.set_xlabel('x'); ax3d_s.set_ylabel('y'); ax3d_s.set_zlabel('z')
        ax_pos_s.plot(df['t'], df['x'], label='x'); ax_pos_s.plot(df['t'], df['y'], label='y'); ax_pos_s.plot(df['t'], df['z'], label='z')
        ax_pos_s.set_title('Position vs time'); ax_pos_s.legend(); ax_pos_s.grid(True)
        ax_ctrl_s.plot(df['t'], df['ax'], label='ax'); ax_ctrl_s.plot(df['t'], df['ay'], label='ay'); ax_ctrl_s.plot(df['t'], df['az'], label='az')
        ax_ctrl_s.set_title('Control vs time'); ax_ctrl_s.legend(); ax_ctrl_s.grid(True)
        ax_info_s.axis('off')
        figs.tight_layout()
        figs.savefig(png_path, dpi=160, bbox_inches='tight')
    print("Saved combined figure to:", png_path)
except Exception as e:
    print("Failed to save combined figure:", e)

if interactive_ok:
    try:
        print("Interactive: close figure window to finish.")
        plt.ioff(); plt.show()
    except Exception:
        pass

# End of file
